import torch
from transformers import BertTokenizer

def predict_sentiment(text, model):
    """
    输入一条文本，输出情感分析类别。
    """
    # 分词和编码
    token = BertTokenizer.from_pretrained('bert-base-chinese')
    encoded = token.batch_encode_plus(
        [text],
        truncation=True,
        padding='max_length',
        max_length=216,
        return_tensors='pt',
        return_length=True,
    )
    input_ids = encoded['input_ids']
    attention_mask = encoded['attention_mask']
    token_type_ids = encoded['token_type_ids']
    
    # 推理
    model.eval()
    with torch.no_grad():
        logits = model.model(input_ids, attention_mask, token_type_ids)
        pred = torch.argmax(logits, dim=1).item()
    
    # 标签映射
    label_map = {0: '负面', 1: '中性', 2: '正面'}
    return label_map.get(pred, str(pred))


if __name__ == '__main__':
    # 示例文本
    example_text = '新冠疫情防控取得了阶段性胜利，大家都很开心！'
    print(f"示例文本: {example_text}")
    print(f"情感分析结果: {predict_sentiment(example_text, model)}")
    example_text = '新冠疫情防控取得了阶段性胜利，大家都很开心！'
    print(f"示例文本: {example_text}")
    print(f"情感分析结果: {predict_sentiment(example_text, model)}")

def analyze_sentiment(text, model=None):
    """
    输入一条文本，输出情感分析类别。
    可以在Notebook的任何cell中直接调用此函数。
    """
    return predict_sentiment(text, model)
# 疫情相关文本列表
texts = [
    "看到新增病例数不断下降，心里真的很开心，感觉胜利在望了！",
    "志愿者们不辞辛劳帮忙运送物资，太让人感动了。",
    "医护人员日夜坚守，为我们的健康保驾护航，致敬！",
    "疫情期间线上办公软件很给力，工作没耽误。",
    "小区防疫措施严格，住着挺安心。",
    "疫情对生活有点影响，但大家都在努力适应。",
    "线上买菜能满足日常需求，就是配送时间有时不太准。",
    "在家隔离不能出门，生活节奏被打乱了。",
    "被隔离期间，工作人员态度不好，沟通很不顺畅。",
    "疫情导致生意冷清，收入大幅减少，压力巨大。"
]

# 分析每个文本的情感
for text in texts:
    sentiment = analyze_sentiment(text, model)